#!/usr/bin/env python3
"""
交互式沙发咨询对话测试工具
支持多轮对话、实时流式输出、状态监控等功能
"""

import sys
import time
import json
from typing import Dict, List, Any

from srd.agents.conversation_agent import SofaConversationAgent

class Colors:
    """简单的颜色控制类"""
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
    GRAY = '\033[90m'

class InteractiveChatTester:
    """交互式对话测试器"""
    
    def __init__(self, table_name: str = "sofa_demo_v2", topk: int = 5):
        self.agent = SofaConversationAgent(table_name=table_name, topk=topk)
        self.conversation_history = []
        self.session_stats = {
            'total_turns': 0,
            'intent_distribution': {},
            'products_recommended': 0,
            'avg_response_time': 0.0
        }
        
    def print_banner(self):
        """打印欢迎横幅"""
        banner = f"""
{Colors.OKCYAN}{'='*80}
{Colors.WARNING}🛋️  产品咨询推荐 AI 助手 - 交互式测试工具
{Colors.OKCYAN}{'='*80}

{Colors.OKGREEN}🚀 功能特性:
{Colors.ENDC}   • 多轮对话支持，保持上下文
   • 实时流式响应输出
   • 智能意图识别 (推荐/详情/聊天/其他)
   • 多模态融合产品检索

{Colors.OKGREEN}📋 可用命令:
{Colors.ENDC}   
   • /clear    - 清理对话历史
   • /stats    - 显示对话统计
   • /exit     - 退出程序

{Colors.OKGREEN}💡 使用提示:
{Colors.ENDC}   • 直接输入问题开始对话
   • 支持产品推荐: "推荐现代风格沙发"
   • 支持详情查询: "第一个产品的材质如何"
   • 支持普通聊天: "你好"

{Colors.OKCYAN}{'='*80}{Colors.ENDC}
"""
        print(banner)
    
    def print_help(self):
        """打印帮助信息"""
        help_text = f"""
{Colors.WARNING}📖 详细使用指南:

{Colors.OKGREEN}🛋️ 产品推荐类问题:{Colors.ENDC}
   • "推荐一些现代简约风格的沙发"
   • "我想要布艺材质的沙发"
   • "预算1万元左右的真皮沙发"
   • "三人位客厅沙发有什么推荐"

{Colors.OKGREEN}🔍 产品详情类问题:{Colors.ENDC}
   • "第一个产品的详细信息"
   • "这款沙发的材质怎么样"
   • "ID为2的产品尺寸规格"
   • "现代简约真皮沙发的保养方法"

{Colors.OKGREEN}💬 普通聊天:{Colors.ENDC}
   • "你好，今天天气不错"
   • "你是什么AI助手"
   • "可以介绍一下你的功能吗"

{Colors.OKGREEN}⚡ 快捷命令:{Colors.ENDC}
   • /clear - 重置对话，开始新的会话
   • /stats - 查看本次会话的统计信息
   • /exit  - 退出测试工具
"""
        print(help_text)
    
    def print_stats(self):
        """打印对话统计"""
        stats = self.session_stats
        intent_summary = ", ".join([f"{k}: {v}" for k, v in stats['intent_distribution'].items()])
        
        stats_text = f"""
{Colors.WARNING}📊 本次会话统计:
{Colors.ENDC}   对话轮次: {stats['total_turns']}
   意图分布: {intent_summary or '无'}
   推荐产品数: {stats['products_recommended']}
   平均响应时间: {stats['avg_response_time']:.2f}秒
"""
        print(stats_text)
    
    def update_stats(self, intent: str, response_time: float, product_count: int = 0):
        """更新统计信息"""
        self.session_stats['total_turns'] += 1
        self.session_stats['intent_distribution'][intent] = \
            self.session_stats['intent_distribution'].get(intent, 0) + 1
        self.session_stats['products_recommended'] += product_count
        
        # 计算平均响应时间
        total_time = self.session_stats['avg_response_time'] * (self.session_stats['total_turns'] - 1)
        self.session_stats['avg_response_time'] = (total_time + response_time) / self.session_stats['total_turns']
    
    def process_user_input(self, user_input: str) -> bool:
        """处理用户输入，返回是否继续"""
        # 处理命令
        if user_input.startswith('/'):
            command = user_input.lower().strip()
            
            if command == '/help':
                self.print_help()
                return True
            elif command == '/clear':
                self.agent.clear_history()
                self.conversation_history.clear()
                self.session_stats = {
                    'total_turns': 0,
                    'intent_distribution': {},
                    'products_recommended': 0,
                    'avg_response_time': 0.0
                }
                print(f"{Colors.OKGREEN}🧹 对话历史已清理，开始新的会话{Colors.ENDC}")
                return True
            elif command == '/stats':
                self.print_stats()
                return True
            elif command in ['/exit', '/quit']:
                print(f"{Colors.WARNING}👋 感谢使用沙发咨询AI助手，再见！{Colors.ENDC}")
                return False
            else:
                print(f"{Colors.FAIL}❌ 未知命令: {command}，输入 /help 查看帮助{Colors.ENDC}")
                return True
        
        # 处理正常对话
        return self.chat_with_agent(user_input)
    
    def chat_with_agent(self, user_input: str) -> bool:
        """与代理进行对话"""
        start_time = time.time()
        
        # 显示用户输入
        print(f"\n{Colors.OKBLUE}👤 用户:{Colors.ENDC} {user_input}")
        
        # 记录到历史
        self.conversation_history.append({"role": "user", "content": user_input})
        
        # 流式处理响应
        intent_detected = None
        products_count = 0
        response_content = []
        
        print(f"{Colors.OKGREEN}🤖 AI助手:{Colors.ENDC} ", end="", flush=True)
        
        try:
            for chunk in self.agent.chat_stream(user_input):
                if chunk['type'] == 'intent':
                    intent_detected = chunk['content']
                    # 显示意图检测结果
                    intent_color = {
                        'product_recommendation': Colors.HEADER,
                        'product_detail_inquiry': Colors.OKCYAN,
                        'normal_chat': Colors.WARNING,
                        'other': Colors.ENDC
                    }.get(intent_detected, Colors.ENDC)
                    
                    print(f"\n{Colors.GRAY}   [意图: {intent_color}{intent_detected}{Colors.GRAY}]{Colors.ENDC}")
                    print(f"{Colors.OKGREEN}🤖 AI助手:{Colors.ENDC} ", end="", flush=True)
                
                elif chunk['type'] == 'products':
                    products_count = len(chunk['content'])
                    print(f"\n{Colors.GRAY}   [推荐了 {products_count} 个产品]{Colors.ENDC}")
                    print(f"{Colors.OKGREEN}🤖 AI助手:{Colors.ENDC} ", end="", flush=True)
                
                elif chunk['type'] == 'content':
                    content_char = chunk['content']
                    print(content_char, end="", flush=True)
                    response_content.append(content_char)
                    time.sleep(0.005)  # 稍微减慢输出速度以便观察
        
        except KeyboardInterrupt:
            print(f"\n{Colors.WARNING}⚠️  响应被用户中断{Colors.ENDC}")
            return True
        except Exception as e:
            print(f"\n{Colors.FAIL}❌ 对话异常: {e}{Colors.ENDC}")
            return True
        
        print()  # 换行
        
        # 记录AI响应到历史
        ai_response = "".join(response_content)
        if ai_response:
            self.conversation_history.append({"role": "assistant", "content": ai_response})
        
        # 计算响应时间
        response_time = time.time() - start_time
        
        # 更新统计
        if intent_detected:
            self.update_stats(intent_detected, response_time, products_count)
        
        # 显示响应时间
        print(f"{Colors.GRAY}   [响应时间: {response_time:.2f}秒]{Colors.ENDC}")
        
        return True
    
    def run(self):
        """运行交互式测试"""
        self.print_banner()
        
        try:
            while True:
                # 显示提示符
                prompt = f"\n{Colors.OKGREEN}💬 请输入您的问题 (或 /help 查看帮助):{Colors.ENDC} "
                user_input = input(prompt).strip()
                
                # 检查空输入
                if not user_input:
                    continue
                
                # 处理输入
                if not self.process_user_input(user_input):
                    break
                    
        except KeyboardInterrupt:
            print(f"\n\n{Colors.WARNING}👋 程序被用户中断，再见！{Colors.ENDC}")
        except EOFError:
            print(f"\n\n{Colors.WARNING}👋 输入结束，再见！{Colors.ENDC}")


def main():
    """主函数"""
    # 检查命令行参数
    table_name = "sofa_demo_v2"
    topk = 5
    
    if len(sys.argv) > 1:
        table_name = sys.argv[1]
    if len(sys.argv) > 2:
        topk = int(sys.argv[2])
    
    # 创建并运行测试器
    tester = InteractiveChatTester(table_name=table_name, topk=topk)
    tester.run()


if __name__ == "__main__":
    main()
